import mindspore
import mindspore.common.dtype as mstype
import mindspore.context as context
import mindspore.nn as nn
import mindspore.ops as ops
import numpy as np
from mindspore.common import initializer as init

from mindspore.ops import functional as F

context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")

class LinearAC(nn.Cell):
    """
    一个AC神经网络的简单示意.
    """
    def __init__(self,
                 input_dense,
                 output_dense,
                 hidden_dense=256,
                 ):
        super(LinearAC, self).__init__()
        self.h1_a = nn.Dense(input_dense, hidden_dense)
        self.h2_a = nn.Dense(hidden_dense, hidden_dense)

        self.h1_c = nn.Dense(input_dense, hidden_dense)
        self.h2_c = nn.Dense(hidden_dense, hidden_dense)

        self.h3 = nn.Dense(hidden_dense, output_dense)
        self.h4 = nn.Dense(hidden_dense, 1)

        self.relu = ops.ReLU()

    def construct(self, x):
        '''构造Actor与Critc同时输出'''
        a_out = self.relu(self.h1_a(x))
        a_out = self.relu(self.h2_a(a_out))
        a_out = ops.Softmax(axis=1)(self.h3(a_out))

        c_out = self.relu(self.h1_c(x))
        c_out = self.relu(self.h2_c(c_out))
        c_out = self.h4(c_out)


        return a_out, c_out
    
class Agent():
    def __init__(self, ac_model:nn.Cell,
        lr = 2e-4,          # 学习率
        p_c = 1.0,          # Critic学习参数
        gamma = 0.98,       # 奖励衰减系数
        lambda_ = 1.0,      # GAE系数
        beta = 0.01,        # 熵损失系数
        epsilon = 0.2,      # PPO裁剪系数
        batch_size = 16,    # 训练批量
        exp_reuse_rate = 10,   # 经验回放复用率
        ) -> None:
        self.ac_model = ac_model
        self.lr = lr
        self.p_c = p_c
        self.gamma = gamma
        self.lambda_ = lambda_     
        self.beta = beta       
        self.epsilon = epsilon      
        self.batch_size = batch_size    
        self.exp_reuse_rate = exp_reuse_rate

    def desision(self, state):
        state = mindspore.Tensor(state).unsqueeze(0).astype(np.float32)

    
    def learn(self, replaybuffer):
        ...


class ExperimentReplayBuffer():
    def __init__(self) -> None:
        pass


if __name__ == "__main__":
    model = LinearAC(4, 2)
    input_ = np.random.randn(4, 4)
    input_ = mindspore.Tensor(input_).astype(np.float32)
    print(model(input_))